'''
Author: xushaocong
Date: 2022-07-03 21:49:53
LastEditTime: 2022-07-03 23:07:18
LastEditors: xushaocong
Description: 
FilePath: /tmp/clf_exp.py
email: xushaocong@stu.xmu.edu.cn
'''


import torch
from loguru import logger

import matplotlib.pyplot as plt
import os 
from tqdm import tqdm
'''
description:  定义的model
return {*}
'''
class Net(torch.nn.Module):  
    def __init__(self, n_feature, n_hidden, n_output):  
        super(Net, self).__init__()  
        self.model =torch.nn.Sequential(
        torch.nn.Linear(n_feature, n_hidden),
        torch.nn.Sigmoid(),
        torch.nn.Linear(n_hidden, n_hidden),
        torch.nn.Sigmoid(),
        torch.nn.Linear(n_hidden, n_output)
        )
        
         
        
    def forward(self, x):  
        return self.model(x)
        
        

'''
description: 保存 训练过程的可视化图像 
param {*} x
param {*} y_red
param {*} save_fig_name
return {*}
'''
def vis(x, y_red,save_fig_name):
    # train result  
    train_predict = torch.max(y_red, 1)[1]  
    plt.scatter(x.data.numpy()[:, 0], x.data.numpy()[:, 1], c=train_predict.data.numpy(), s=100, lw=0, cmap='RdYlGn')  
    plt.savefig(save_fig_name) 


def main():
    sample_num  = 100#* 标签为0 的100个 , 标签为1的100 个
    epoch = 500 
    vis_fig_name="vis_%05d.jpg"
    #* 初始化数据集
    n_data = torch.ones(sample_num, 2)  
    x0 = torch.normal(2*n_data, 1) #* 生成均值为2 , 反差为1 的 (100,2) 的 的数据集
    y0 = torch.zeros(sample_num)  
    x1 = torch.normal(-2*n_data, 1)#* 生成均值为-2 , 反差为1 的 (100,2) 的 的数据集  
    y1 = torch.ones(sample_num)  
    #* 拼接
    x = torch.cat((x0, x1)).type(torch.FloatTensor)  
    y = torch.cat((y0, y1)).type(torch.LongTensor) 

    net = Net(n_feature=2, n_hidden=10, n_output=2)
    optimizer = torch.optim.Adam(net.parameters(), lr=0.02)  
    loss_func = torch.nn.CrossEntropyLoss() 


    net = net.cuda()
    losses = []
    for i in tqdm(range(epoch)):  
        x = x.cuda()
        y = y.cuda()
        out = net(x)   
        loss = loss_func(out, y)  
        optimizer.zero_grad()  
        loss.backward()  
        optimizer.step() 
        logger.info(f"loss == {loss}")
        losses.append(loss)
        vis(x.cpu(),out.cpu(),vis_fig_name%i)#* 可视化每个epoch的量化结果

    logger.info("train finish ")
    train_result = net(x)  
    vis(x.cpu(),train_result.cpu(),vis_fig_name%(i+1))


    
if __name__ =="__main__":
    os.environ["VISIBLE_CUDA_DEVICES"] = "3"
    main()
